import asyncio
import os
import dotenv
import random

from openai import AsyncOpenAI
from openai.types.responses import ResponseTextDeltaEvent

from agents import (
    Agent,
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    RunConfig,
    Runner,
    function_tool,
    set_tracing_disabled,
    ItemHelpers,
)

dotenv.load_dotenv()

client = AsyncOpenAI(api_key=os.getenv("AI_API_KEY"),base_url=os.getenv("AI_BASE_URL"))
set_tracing_disabled(disabled=True)

class CustomModelProvider(ModelProvider):
    def get_model(self, model_name: str | None) -> Model:
        return OpenAIChatCompletionsModel(model=os.getenv("AI_API_MODEL"), openai_client=client)

CUSTOM_MODEL_PROVIDER = CustomModelProvider()


@function_tool
def howTool()->int:
    r = random.randint(1, 10)
    print(f"howTool返回了{r}")
    return r


async def main():
    agent = Agent(
        name="王富贵",
        instructions="""
        你必须先调用 howTool 工具，获取一个 1~10 的随机整数 N。然后你要根据 N，连续输出 N 个不同的笑话，每个笑话单独成段。
        例如：如果 howTool 返回 3，你就输出 3 个不同的笑话，每个笑话单独成段。
        """,
        tools=[howTool],
    )
    result = Runner.run_streamed(agent, input="开始", run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),)

    # result2 = await Runner.run(
    #     agent,
    #     input="开始",
    #     run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),
    # )
    # print(result2.final_output)

    last_message = ""
    async for event in result.stream_events():
        if event.type == "raw_response_event":
            print(f"原始响应事件: {event.type}")
            continue
        elif event.type == "agent_updated_stream_event":
            print(f"Agent已更新: {event.new_agent.name}")
            continue
        elif event.type == "run_item_stream_event":
            if event.item.type == "tool_call_item":
                print("-- 工具被调用")
            elif event.item.type == "tool_call_output_item":
                print(f"-- 工具输出: {event.item.output}")
            elif event.item.type == "message_output_item":
                msg = ItemHelpers.text_message_output(event.item)
                # print(f"-- 消息输出原始内容: {event.item}")
                if msg.strip():
                    last_message = msg
            else:
                print(f"-- 其他事件类型: {event.item.type}")
                pass
    if last_message:
        print("-- 最终笑话输出:\n" + last_message)

if __name__ == "__main__":
    asyncio.run(main())



